#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include "types.h"
#include "elf_types.h"
#include "err.h"

#define ELF_BUF_SIZ 0x40000U

static int check_ehdr (Elf32_Ehdr * ehdr);

byte elf_buf[ELF_BUF_SIZ];

// int
// main (int argc, char ** argv) {
//     int elf_fd;
//     int seg_num;
//     Elf32_Ehdr ehdr_buf;
//     Elf32_Phdr phdr_buf;
//     if (argc != 2) {
//         printf ("Usage: %s elf_file.\n", argv[0]);
//         return 0;
//     }
//     if ((elf_fd = open (argv[1], O_RDONLY)) == -1) {
//         perror ("Unable to open.\n");
//     }
//     if (read (elf_fd, &ehdr_buf, sizeof (Elf32_Ehdr)) != sizeof (Elf32_Ehdr)) {
//         perror ("Unable to read Ehdr.\n");
//     }
//     printf ("phnum:%d\nphoff:%d\n", ehdr_buf.e_phnum, ehdr_buf.e_phoff);
//     seg_num = ehdr_buf.e_phnum;
//     lseek (elf_fd, ehdr_buf.e_phoff, SEEK_SET);
//     if (read (elf_fd, &phdr_buf, sizeof (Elf32_Phdr)) != sizeof (Elf32_Phdr)) {
//         perror ("Unable to read phdr.\n");
//     } else {
//         printf ("f_off:%lX -> v_addr:%X, %d byte. type:%X, flag:%X\n",phdr_buf.p_offset + sizeof(Elf32_Ehdr) + sizeof (Elf32_Phdr) * ehdr_buf.e_phnum, phdr_buf.p_vaddr, phdr_buf.p_filesz , phdr_buf.p_type, phdr_buf.p_flags);
//     }
//     if (read (elf_fd, &phdr_buf, sizeof (Elf32_Phdr)) != sizeof (Elf32_Phdr)) {
//         perror ("Unable to read phdr.\n");
//     } else {
//         printf ("f_off:%lX -> v_addr:%X, %d byte. type:%X, flag:%X\n",phdr_buf.p_offset + sizeof(Elf32_Ehdr) + sizeof (Elf32_Phdr) * ehdr_buf.e_phnum ,phdr_buf.p_vaddr, phdr_buf.p_filesz , phdr_buf.p_type, phdr_buf.p_flags);
//     }
// }

int
elf32_load_kernel (const char* kernel_path, byte * mem, uint32 max_mem) {
    int stat;
    int elf_fd, n, i, j, seg_num;
    uint32 seg_base, seg_hdr_base, filesz, memsz, virt_base, virt_off, rdnum, rdoff;
    Elf32_Ehdr ehdr_buf;
    Elf32_Phdr phdr_buf;
    stat = 0;
    if ((elf_fd = open (kernel_path, O_RDONLY)) == -1) {
        perror ("Unable to open kernel file.\n");
        stat |= E_ELF_FILE_ERR;
        goto error;
    }
    #define EHDR_SIZ (sizeof (Elf32_Ehdr))
    // #define PHDR_SIZ (sizeof (Elf32_Phdr))
    #define PHDR_SIZ (ehdr_buf.e_phentsize)
    if ((n = read (elf_fd, &ehdr_buf, EHDR_SIZ)) != EHDR_SIZ) {
        perror ("Invalid kernel file (too small).");
        stat |= E_ELF_FILE_ERR;
        goto error;
    }
    if (check_ehdr (&ehdr_buf) != 0) {
        fprintf (stderr, "Invalid kernel file (elf header check failed).");
        stat |= E_ELF_FILE_ERR;
        goto error;
    }
    seg_base = ehdr_buf.e_phoff + ehdr_buf.e_phnum * PHDR_SIZ;
    seg_hdr_base = ehdr_buf.e_phoff;
    seg_num = ehdr_buf.e_phnum;
    printf ("seg_num:%d\n", seg_num);
    for (i = 0; i < seg_num; i++) {
        printf ("loading seg %d\n", i);
        if (lseek (elf_fd, seg_hdr_base + i * PHDR_SIZ, SEEK_SET) == -1) {
            perror ("Unable to seek.");
            stat |= E_ELF_FILE_ERR;
            goto error;
        }
        if ((n = read (elf_fd, &phdr_buf, PHDR_SIZ)) != PHDR_SIZ) {
            perror ("Invalid kernel file (too small).");
            stat |= E_ELF_FILE_ERR;
            goto error;
        }
        if (phdr_buf.p_type != PT_LOAD) continue; 
        filesz = phdr_buf.p_filesz;
        memsz = phdr_buf.p_memsz;
        virt_base = phdr_buf.p_vaddr;
        virt_off = virt_base;
        if (virt_base + memsz >= max_mem) {
            fprintf (stderr, "Invalid kernel file (insuffcient memory).\n");
            stat |= E_ELF_NEED_MORE_MEM;
            goto error;
        }
        rdoff = seg_base + phdr_buf.p_offset;
        if (lseek (elf_fd, rdoff, SEEK_SET) == -1) {
            perror ("Unable to seek.");
            stat |= E_ELF_FILE_ERR;
            goto error;
        }
        if (filesz == 0) printf ("seg%d : Empty segment.\n", i);
        memset (mem + virt_off, 0, memsz);
        while (filesz > 0) {
            if (filesz < ELF_BUF_SIZ) {
                rdnum = filesz;
            } else {
                rdnum = ELF_BUF_SIZ;
            }
            if ((n = read (elf_fd, elf_buf, rdnum)) != rdnum) {
                fprintf (stderr, "Invalid kernel file (too small).\n");
                stat |= E_ELF_FILE_ERR;
                goto error;
            }
            memcpy (mem + virt_off, elf_buf, rdnum);
            printf ("seg%d : %d bytes -> 0x%X\n", i, rdnum, virt_off);
            virt_off += rdnum;
            filesz -= rdnum;
        }
    }
    success:
    return stat;
    error:
    return stat;
}

static int
check_ehdr (Elf32_Ehdr * ehdr) {
    if (ehdr->e_ident[0] == 0x7f)
    if (ehdr->e_ident[1] == 'E')
    if (ehdr->e_ident[2] == 'L')
    if (ehdr->e_ident[3] == 'F')
    if (ehdr->e_ident[4] == 1) // 32
    if (ehdr->e_ident[5] == 1) // lsb
    if (ehdr->e_types == 2) // exec
    if (ehdr->e_machine == 243) // riscv
    return 0;
    return -1;
}